﻿using StableDiffusionTools.Oobabooga;
using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Text;
using System.Text.Json.Serialization;
using System.Threading.Tasks;

namespace LLMEvaluation.Framework
{
    public class RunEntry : INotifyPropertyChanged
    {
        public List<Message> PromptMessages { get; set; }
        public List<ModelEntry> ModelEntries { get; set; }
        public int Seed { get; set; } = -1;
        public int Runs { get; set; } = 1;

        private string prompt = "";
        [JsonIgnore(Condition = JsonIgnoreCondition.Always)]
        public string Prompt
        {
            get => prompt;
            set
            {
                prompt = value;
                NotifyPropertyChanged();
            }
        }

        private string models = "";
        [JsonIgnore(Condition = JsonIgnoreCondition.Always)]
        public string Models
        {
            get => models;
            set
            {
                models = value;
                NotifyPropertyChanged();
            }
        }

        private InferenceState inference = InferenceState.Idle;
        [JsonIgnore(Condition = JsonIgnoreCondition.Always)]
        public InferenceState Inference
        {
            get => inference;
            set
            {
                inference = value;
                NotifyPropertyChanged();
            }
        }

        public event PropertyChangedEventHandler PropertyChanged;
        private void NotifyPropertyChanged([CallerMemberName] string propertyName = "")
        {
            PropertyChanged?.Invoke(this, new PropertyChangedEventArgs(propertyName));
        }

        public RunEntry()
        {
            UpdateAll();
        }

        public Message AddPromptMessage(Role role, string content)
        {
            return AddPromptMessage(new Message(role, content));
        }

        public Message AddPromptMessage(Message message)
        {
            PromptMessages.Add(message);
            UpdatePromptMessages();
            return message;
        }

        public void AddPromptMessages(List<Message> messages, int insertIndex = -1)
        {
            if (insertIndex == -1)
                PromptMessages.AddRange(messages);
            else
                PromptMessages.InsertRange(insertIndex, messages);

            UpdatePromptMessages();
        }

        public void RemovePromptMessage(Message message)
        {
            if (PromptMessages.Remove(message))
                UpdatePromptMessages();
        }

        public void RemovePromptMessages(IList<Message> messages)
        {
            for (int i = messages.Count - 1; i >= 0; i--)
                for (int j = PromptMessages.Count - 1; j >= 0; j--)
                    if (messages[i] == PromptMessages[j])
                        PromptMessages.RemoveAt(j);

            UpdatePromptMessages();
        }

        public void UpdatePromptMessages()
        {
            if (PromptMessages == null)
                PromptMessages = new List<Message>();

            if (PromptMessages.Count == 0)
                Prompt = "None";
            else
            {
                StringBuilder sb = new StringBuilder(string.Join(", ", PromptMessages));
                sb.Replace("\r\n", " ").Replace("\n", " ");
                Prompt = sb.ToString();
            }

            NotifyPropertyChanged("PromptMessages");
        }

        public ModelEntry AddModel(string name)
        {
            return AddModel(new ModelEntry() { Name = name });
        }

        public ModelEntry AddModel(ModelEntry model)
        {
            ModelEntries.Add(model);
            UpdateModels();
            return model;
        }

        public void AddModels(List<ModelEntry> models, int insertIndex = -1)
        {
            if (insertIndex == -1)
                ModelEntries.AddRange(models);
            else
                ModelEntries.InsertRange(insertIndex, models);

            UpdateModels();
        }

        public void RemoveModel(ModelEntry model)
        {
            if (ModelEntries.Remove(model))
                UpdateModels();
        }

        public void RemoveModels(IList<ModelEntry> models)
        {
            for (int i = models.Count - 1; i >= 0; i--)
                for (int j = ModelEntries.Count - 1; j >= 0; j--)
                    if (models[i] == ModelEntries[j])
                        ModelEntries.RemoveAt(j);

            UpdateModels();
        }

        public void UpdateModels()
        {
            if (ModelEntries == null)
                ModelEntries = new List<ModelEntry>();

            Models = string.Join(", ", ModelEntries);
        }


        public void UpdateAll()
        {
            UpdatePromptMessages();
            UpdateModels();
            for (int i = 0; i < ModelEntries.Count; i++)
            {
                ModelEntries[i].Update();
            }
        }
    }
}
